// New high-performance, lightweight, non-blocking network server framework written in Go.
// Description: A simple TCP server that receives data from DTU devices and saves it to MySQL(influxdb).
// Author: zengkai001@qq.com, 2025.03.18

// # In project directory:             cd /home/kzeng/pss_dtu
// # Reset environment:                go clean -modcache
//                                     go env -w GOPROXY=https://goproxy.cn,direct

// # Initialize module (if missing):   go mod init pss_dtu
// # Install dependencies:             go mod tidy
// Install gnet:                       go get -u github.com/panjf2000/gnet/v2
// Install MySQL connect driver:	   go get -u github.com/go-sql-driver/mysql
// Install Influxdb connect driver:    go get github.com/influxdata/influxdb-client-go/v2

// run: 							   go run main_influxdb.go

// test:                               nc <ip_address> 55555

package main

import (
	"log"
	"strings"
	"time"
	"os"
	"os/signal"
	"syscall"

	"github.com/panjf2000/gnet/v2"
	influxdb2 "github.com/influxdata/influxdb-client-go/v2"
	"github.com/influxdata/influxdb-client-go/v2/api"
)

// Save data to InfluxDB
const SAVE_INFLUXDB = true

// MyTCPServer
type MyTCPServer struct {
	gnet.BuiltinEventEngine
	influxClient influxdb2.Client
	writeAPI     api.WriteAPI
	errorsCh     chan error
}

// OnBoot: handle server start
func (ts *MyTCPServer) OnBoot(eng gnet.Engine) (action gnet.Action) {
	log.Println("TCP Server started on port 55555")
	if SAVE_INFLUXDB {
		log.Println("Data will be saved to InfluxDB.")
		// Start error handling goroutine
		go ts.handleErrors()
	}
	return
}

// OnOpen: handle new connection
func (ts *MyTCPServer) OnOpen(c gnet.Conn) (out []byte, action gnet.Action) {
	log.Printf("New connection from %s\n", c.RemoteAddr().String())
	return
}

// OnTraffic: handle data traffic
func (ts *MyTCPServer) OnTraffic(c gnet.Conn) (action gnet.Action) {
	buffer, err := c.Next(-1) // read all data
	if err != nil {
		log.Printf("Error reading data: %v\n", err)
		return
	}

	if len(buffer) == 0 {
		return
	}

	// parse data
	data := strings.TrimSpace(string(buffer))
	log.Printf("--->>>Received data: '%s' from %s\n", data, c.RemoteAddr().String())

	dtuID, dtuDataMap := parseDTUData(data)
	if errMsg, exists := dtuDataMap["error"]; exists {
		log.Printf("Failed to parse data: %s\n", errMsg)
		return
	}
	log.Printf("Parsed data: dtu_id=%s, dtu_data=%v\n", dtuID, dtuDataMap)

	// if SAVE_INFLUXDB is true, save data to InfluxDB
	if SAVE_INFLUXDB {
		go func() { // 使用 goroutine 异步保存数据
			err := ts.saveToInfluxDB(dtuID, dtuDataMap)
			if err != nil {
				log.Printf("Failed to save data to InfluxDB: %v\n", err)
			} else {
				log.Printf("--->>>Data saved to InfluxDB: dtu_id=%s, dtu_data=%v\n", dtuID, dtuDataMap)
			}
		}()
	}

	return
}

// OnClose: handle connection close
func (ts *MyTCPServer) OnClose(c gnet.Conn, err error) (action gnet.Action) {
	log.Printf("Connection closed: %s\n", c.RemoteAddr().String())
	return
}

// parseDTUData: Parse a comma-separated string
func parseDTUData(rawData string) (string, map[string]string) {
	// Split the raw data by commas
	log.Printf("Raw data: %s\n", rawData)

	parts := strings.Split(rawData, ",")
	if len(parts) < 12 {
		return "UNKNOWN", map[string]string{"error": "Invalid data format"}
	}

	// Extract dtu_id and other fields
	dtuID := parts[0]
	if len(dtuID) > 15 {
		dtuID = dtuID[:15]
	}

	dataMap := map[string]string{
		"temp_001":            parts[1],
		"vib_001":             parts[2],
		"cur_001":             parts[3],
		"press_01":            parts[4],
		"spd_001":             parts[5],
		"status":              parts[6],
		"err_1":               parts[7],
		"err_2":               parts[8],
		"env_temp_room1":      parts[9],
		"env_humidity_room1":  parts[10],
		"alarm":               parts[11],
	}

	return dtuID, dataMap
}

// Save data to InfluxDB
var shanghaiLoc = time.FixedZone("CST", 8*60*60) // UTC+8

// Add error handling
func (ts *MyTCPServer) handleErrors() {
	for err := range ts.errorsCh {
		log.Printf("Write Error: %s\n", err.Error())
	}
}

// Modify saveToInfluxDB to use batch processing
func (ts *MyTCPServer) saveToInfluxDB(dtuID string, dtuData map[string]string) error {
	point := influxdb2.NewPointWithMeasurement("device_data").
		AddTag("dtu_id", dtuID)
	
	// Add fields individually with type conversion
	for key, value := range dtuData {
		if key != "error" {
			point.AddField(key, value)
		}
	}
	
	point.SetTime(time.Now().In(shanghaiLoc))
	ts.writeAPI.WritePoint(point)
	return nil
}

func main() {
	// Initialize InfluxDB client with options
	options := influxdb2.DefaultOptions().
		SetBatchSize(100).
		SetFlushInterval(5000).
		SetRetryInterval(1000).
		SetMaxRetries(3).
		SetMaxRetryTime(10000)

	// Kzeng influxDB Info
	// influxClient := influxdb2.NewClientWithOptions(
	// 	"http://localhost:8686",
	// 	"q8PJ4wDYJJ0KHe60fFIBIB70pL2JlSZVpvDf4rz7-PiDQDasC38uLBayqJ9_iRgnAgG210Vq2XsVyNIprrunWw==",
	// 	options,
	// )
	// writeAPI := influxClient.WriteAPI("yc-tech", "ddb")


	// YC-TECH influxDB Info
	influxClient := influxdb2.NewClientWithOptions(
		"http://localhost:8086",
		"hp4h6Dh8ilTlc4V6lyC1Nl_PFhd93Se9aJtaJFFaWNI-Slg864yCCRBbjIjTCQUqUWXRfsSpsdWayEOL1a6a_g==",
		options,
	)
	writeAPI := influxClient.WriteAPI("default", "ddb")


	errorsCh := make(chan error)

	// Correctly set the write failed callback
	go func() {
		for err := range writeAPI.Errors() {
			errorsCh <- err
		}
	}()

	// Create signal channel for graceful shutdown
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	// Start TCP server
	tcpServer := &MyTCPServer{
		influxClient: influxClient,
		writeAPI:     writeAPI,
		errorsCh:     errorsCh,
	}

	// Handle shutdown gracefully
	go func() {
		<-sigCh
		log.Println("Shutting down server...")
		writeAPI.Flush()
		influxClient.Close()
		close(errorsCh)
		os.Exit(0)
	}()

	log.Println("Starting TCP server on port 55555...")
	if err := gnet.Run(tcpServer, "tcp://:55555",
		gnet.WithMulticore(false), // 禁用多核
		gnet.WithReusePort(false), // 禁用端口复用
		gnet.WithTCPKeepAlive(time.Minute),
	); err != nil {
		log.Fatalf("TCP Server failed: %v", err)
	}
}
